#
# firslov 2022.04.30
#

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import numpy as np
import copy


# word Embedding
class Embeddings(nn.Module):
    """
    位于 word2index 之后
    vocab: 词表数量
    d_model: 词嵌入维度

    """

    def __init__(self, vocab, d_model):
        super(Embeddings, self).__init__()

        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)


# 位置编码器
class PositionalEncoding(nn.Module):
    """
    采用正余弦编码
    d_model: 词嵌入维度
    dropout: 未设默认值
    """

    def __init__(self, d_model, dropout, max_len=1000):
        super(PositionalEncoding, self).__init__()

        self.dropout = nn.Dropout(p=dropout)
        # max_len * d_model
        pe = torch.zeros(max_len, d_model)
        # max_len * 1
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2)
                             * -(math.log(10000.0) / d_model))

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)

        # pe不是参数，不做更新
        self.register_buffer('pe', pe)

    def forward(self, x):
        # x为文本序列的词嵌入表示
        # 限制pe第2维度为句子最大长度
        x = x + self.pe[:, :x.size(1)].requires_grad_(False)
        return self.dropout(x)


# 掩码张量
def subsequent_mask(size):
    attn_shape = (1, size, size)
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
    return torch.from_numpy(1 - subsequent_mask)


# Attention
def attention(query, key, value, mask=None, dropout=None):
    """
    q, k, v: 注意力输入张量
    mask: 掩码张量
    dropout: Dropout Object
    """
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)

    p_attn = F.softmax(scores, dim=-1)

    if dropout is not None:
        p_attn = dropout(p_attn)

    return torch.matmul(p_attn, value), p_attn


# 克隆函数
def clones(module, N):
    """
    module: 克隆模块
    N: 克隆数量
    """
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


# 多头注意力模块
class MultiHeadedAttention(nn.Module):
    """
    head: 头数
    embedding_dim: 词嵌入维度
    dropout: 默认0.1
    """

    def __init__(self, head, embedding_dim, dropout=0.1):
        super(MultiHeadedAttention, self).__init__()

        assert embedding_dim % head == 0

        self.d_k = embedding_dim // head
        self.head = head
        self.embedding_dim = embedding_dim
        self.linears = clones(nn.Linear(embedding_dim, embedding_dim), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        if mask is not None:
            mask = mask.unsqueeze(1)

        batch_size = query.size(0)

        query, key, value = \
            [model(x).view(batch_size, -1, self.head, self.d_k).transpose(1, 2)
             for model, x in zip(self.linears, (query, key, value))]

        x, self.attn = attention(
            query, key, value, mask=mask, dropout=self.dropout)
        x = x.transpose(1, 2).contiguous().view(
            batch_size, -1, self.head * self.d_k)

        return self.linears[-1](x)


# 前馈全连接网络
class PositionwiseFeedForward(nn.Module):
    """
    d_model: 词嵌入维度
    d_ff: 隐层节点数
    dropout: 默认0.1
    """

    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()

        self.w1 = nn.Linear(d_model, d_ff)
        self.w2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, x):
        return self.w2(self.dropout(F.relu(self.w1(x))))


# 规范化层的类
class LayerNorm(nn.Module):
    """
    d_model: 词嵌入维度
    """

    def __init__(self, d_model, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a2 = nn.Parameter(torch.ones(d_model))
        self.b2 = nn.Parameter(torch.zeros(d_model))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a2 * (x - mean) / (std + self.eps) + self.b2


# 子层连接结构
class SublayerConnection(nn.Module):
    """
    d_model: 词嵌入维度
    dropout: 默认0.1
    """

    def __init__(self, d_model, dropout=0.1):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(d_model)
        self.dropout = nn.Dropout(p=dropout)
        self.size = d_model

    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x)))


# 编码器层的类
class EncoderLayer(nn.Module):
    """
    d_model: 词嵌入维度
    self_attn: self-Attention Object
    feed_forward: Feedforward Object
    dropout: 无默认值
    """

    def __init__(self, d_model, self_attn, feed_forward, dropout):
        super(EncoderLayer, self).__init__()

        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.size = d_model

        self.sublayer = clones(SublayerConnection(d_model, dropout), 2)

    def forward(self, x, mask):
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)


# 编码器类
class Encoder(nn.Module):
    """
    layer: EncoderLayer
    N: 编码器层数量
    """

    def __init__(self, layer, N):
        super(Encoder, self).__init__()

        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask):
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)


# 解码器层类
class DecoderLayer(nn.Module):
    """
    d_model: 词嵌入维度
    self_attn: self-Attention Object
    src_attn: Attention Object
    feed_forward: Feedforward Object
    dropout: 无默认值
    """

    def __init__(self, d_model, self_attn, src_attn, feed_forward, dropout):
        super(DecoderLayer, self).__init__()

        self.size = d_model
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.dropout = dropout

        self.sublayer = clones(SublayerConnection(d_model, dropout), 3)

    def forward(self, x, memory, source_mask, target_mask):
        """
        memory: 编码器输出
        source_mask: 编码器输出mask
        target_mask: 解码器输入mask
        """
        m = memory
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, target_mask))
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, source_mask))
        return self.sublayer[2](x, self.feed_forward)


# 解码器类
class Decoder(nn.Module):
    """
    layer: DecoderLayer
    N: 数量
    """

    def __init__(self, layer, N):
        super(Decoder, self).__init__()

        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, memory, source_mask, target_mask):
        """
        memory: 编码器输出
        source_mask: 编码器输出mask
        target_mask: 解码器输入mask
        """

        for layer in self.layers:
            x = layer(x, memory, source_mask, target_mask)
        return self.norm(x)


# Generator类
class Generator(nn.Module):
    """
    d_model: 词嵌入维度
    vocab_size: 词表单词数量
    """

    def __init__(self, d_model, vocab_size):
        super(Generator, self).__init__()

        self.project = nn.Linear(d_model, vocab_size)

    def forward(self, x):
        return F.log_softmax(self.project(x), dim=-1)


# 编码器-解码器结构类
class EncoderDecoder(nn.Module):
    """
    encoder: Encoder Object
    decoder: Decoder Object
    source_embed: 词嵌入+位置编码后的原数据
    target_embed: 词嵌入+位置编码后的目标数据
    generator: Generator Object
    """

    def __init__(self, encoder, decoder, source_embed, target_embed, generator):

        super(EncoderDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = source_embed
        self.tgt_embed = target_embed
        self.generator = generator

    def forward(self, source, target, source_mask, target_mask):
        return self.decode(self.encode(source, source_mask), source_mask, target, target_mask)

    def encode(self, source, source_mask):
        return self.encoder(self.src_embed(source), source_mask)

    def decode(self, memory, source_mask, target, target_mask):
        return self.decoder(self.tgt_embed(target), memory, source_mask, target_mask)


# Transformer
class Transformer(EncoderDecoder):
    """
    source_vocab: 源词汇数
    target: 目标词汇数
    N: 编/解码器层数
    d_model: 词嵌入维度
    d_ff: 输出层隐层节点数
    head: 头数
    dropout: 默认0.1
    """

    def __init__(self, source_vocab, target_vocab, N=6, d_model=512, d_ff=2048, head=8, dropout=0.1):
        c = copy.deepcopy
        attn = MultiHeadedAttention(head, d_model)
        ff = PositionwiseFeedForward(d_model, d_ff, dropout)
        position = PositionalEncoding(d_model, dropout)

        super(Transformer, self).__init__(Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
                                          Decoder(DecoderLayer(d_model, c(
                                              attn), c(attn), c(ff), dropout), N),
                                          nn.Sequential(Embeddings(
                                              source_vocab, d_model), c(position)),
                                          nn.Sequential(Embeddings(
                                              target_vocab, d_model), c(position)),
                                          Generator(d_model, target_vocab)
                                          )
        self.init_weight()

    def init_weight(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)
